from asyncio.log import logger
from unicodedata import category
from easydict import EasyDict
import yaml
import glob
from pathlib import Path
import time
import datetime
import numpy as np
import torch
import rospy
from enum import Enum
# from object_detection.msg import box3d
from object_tracking.msg import trackbox
import ros_numpy as rnp
from std_msgs.msg import String
from sensor_msgs.msg import PointCloud2
from visualization_msgs.msg import Marker
from geometry_msgs.msg import Point

from pcdet.utils import common_utils
from pcdet.datasets import DatasetTemplate
from pcdet.models import build_network, load_data_to_gpu
from pcdet.utils.box_utils import boxes_to_corners_3d


def cfg_from_yaml_file(cfg_file, config, base_config_file=None):
    def merge_new_config(config, new_config):
        if '_BASE_CONFIG_' in new_config:
            with open(base_config_file, 'r') as f:
                try:
                    yaml_config = yaml.safe_load(f, Loader=yaml.FullLoader)
                except:
                    yaml_config = yaml.safe_load(f)
            config.update(EasyDict(yaml_config))

        for key, val in new_config.items():
            if not isinstance(val, dict):
                config[key] = val
                continue
            if key not in config:
                config[key] = EasyDict()
            merge_new_config(config[key], val)

    with open(cfg_file, 'r') as f:
        try:
            new_config = yaml.safe_load(f, Loader=yaml.FullLoader)
        except:
            new_config = yaml.safe_load(f)

        merge_new_config(config=config, new_config=new_config)

    return config


def get_xyzi_points(cloud_array, remove_nans=True, dtype=np.float):
    '''Pulls out x, y, and z columns from the cloud recordarray, and returns
	a 3xN matrix.
    '''
    # remove crap points
    if remove_nans:
        mask = np.isfinite(cloud_array['x']) & np.isfinite(cloud_array['y']) & np.isfinite(cloud_array['z']) & np.isfinite(cloud_array['i'])
        cloud_array = cloud_array[mask]
    
    # pull out x, y, and z values
    points = np.zeros(cloud_array.shape + (4,), dtype=dtype)
    points[...,0] = cloud_array['x']
    points[...,1] = cloud_array['y']
    points[...,2] = cloud_array['z']
    points[...,3] = cloud_array['i']

    return points

def pointcloud2_to_xyzi_array(cloud_msg, remove_nans=True):
    return get_xyzi_points(rnp.point_cloud2.pointcloud2_to_array(cloud_msg), remove_nans=remove_nans)



class ClassName(Enum):
    car = 1
    person = 2
    cyclist = 3

def exetime(func):
    def newfunc(*args, **args2):
        _t0  = time.time()
        back = func(*args, **args2)
        _t1= time.time()
        print("{:20s}".format(func.__name__) + " : " + "{:5.1f}".format((_t1 - _t0) * 1000) + "ms")
        return back
    return newfunc

class DemoDataset(DatasetTemplate):
    def __init__(self, dataset_cfg, class_names, training=False, root_path=None, logger=None, ext='.bin'):
        super().__init__(
            dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger
        )

    def __len__(self):
        return 0

    def __getitem__(self, index):
        return dict()

class Detector:
    def __init__(self, input_dict):

        cfg = EasyDict()
        cfg_from_yaml_file(input_dict.model_cfg, cfg, input_dict.dataset_cfg)

        # print(cfg)
        self.demo_dataset = DemoDataset(
            dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES
        )

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=self.demo_dataset)
        self.model.load_params_from_file(filename=input_dict.ckpt_file, to_cpu=False, logger=common_utils.create_logger())
        self.model.to(self.device)
        self.model.cuda()
        self.model.eval()

        #for ROS 
        self.sub = rospy.Subscriber(input_dict.topic, PointCloud2, self.pc2_callback, queue_size=1)
        self.mk_pub = rospy.Publisher("ros_detector", Marker, queue_size=1)
        self.box_pub = rospy.Publisher("track_box", trackbox, queue_size= 10)
        self.frame_id = input_dict.lidar_frame
        self.score_threshold = input_dict.score_threshold
        self.z_offset = input_dict.z_offset
        self.use_intensity = input_dict.use_intensity
        

    @exetime
    def detector(self, points):

        with torch.no_grad():
            input_dict = {
                'points': points,
                'frame_id': self.frame_id,
            }
            data_dict = self.demo_dataset.prepare_data(data_dict=input_dict)
            data_dict = self.demo_dataset.collate_batch([data_dict])
            load_data_to_gpu(data_dict)

            pred_dicts = self.model.forward(data_dict)[0][0]  # batch_size = 1
            # print(pred_dicts)
            # print(pred_dicts['pred_boxes'].shape)
            # box [num, 7]
            # score [num, ]
            pred_labels = pred_dicts['pred_labels'].detach().cpu().numpy()
            pred_boxes = pred_dicts['pred_boxes'].detach().cpu().numpy()
            pred_scores= pred_dicts['pred_scores'].detach().cpu().numpy()
            
            pred_boxes = pred_boxes[pred_scores > self.score_threshold, :]
            
            return pred_boxes


    @exetime
    def pc2_callback(self, msg):
        # 1. get xyz
        if not self.use_intensity:
            points_raw = rnp.point_cloud2.pointcloud2_to_xyz_array(msg)
            points_raw = np.hstack((points_raw, np.zeros([len(points_raw), 1])))
        else:
            points_raw = pointcloud2_to_xyzi_array(msg)

        #the height of velodyne is 1.78m
        # print(np.mean(points_raw[:, 2]))
        # average point cloud height is -1.26, -1.26-0.518 = -1.778
        points_raw[:, 2] = points_raw[:, 2] - self.z_offset
        # points_raw[:, 2] = points_raw[:, 2]
        
        bbox3d = self.detector(points_raw)

        self.pub_trackbox(bbox3d=bbox3d, frame_id=self.frame_id)
        
        
        # self.frame_id += 1
        print('-------------------------------------')


    def viz(self, bbox3d, frame_id):
        marker = Marker()
        marker.header.frame_id = frame_id
        marker.type = marker.LINE_LIST
        marker.action = marker.ADD
        marker.id = 0
        marker.header.stamp = rospy.Time.now()

        # marker scale (scale y and z not used due to being linelist)
        marker.scale.x = 0.08
        # marker color
        marker.color.a = 1.0
        marker.color.r = 1.0
        marker.color.g = 1.0
        marker.color.b = 0.0

        marker.pose.position.x = 0.0
        marker.pose.position.y = 0.0
        marker.pose.position.z = 0.0

        marker.pose.orientation.x = 0.0
        marker.pose.orientation.y = 0.0
        marker.pose.orientation.z = 0.0
        marker.pose.orientation.w = 1.0
        marker.points = []
        corner_for_box_list = [0, 1, 0, 3, 2, 3, 2, 1, 4, 5, 4, 7, 6, 7, 6, 5, 3, 7, 0, 4, 1, 5, 2, 6]
        if bbox3d.shape[0] == 0:
            print("no %s is detected"%self.category)
            print(" ")
        else: 
            corners3d = boxes_to_corners_3d(bbox3d)  # (N,8,3)
            for box_nr in range(corners3d.shape[0]):
                if abs(corners3d[box_nr][0][0]) > 50 or abs(corners3d[box_nr][0][1]) > 50:
                    continue
                box3d_pts_3d_velo = corners3d[box_nr]  # (8,3)
                for corner in corner_for_box_list:
                    transformed_p = np.array(box3d_pts_3d_velo[corner, 0:4])
                    # transformed_p = transform_point(p, np.linalg.inv(self.Tr_velo_kitti_cam))
                    p = Point()
                    p.x = transformed_p[0]
                    p.y = transformed_p[1]
                    p.z = transformed_p[2] + self.z_offset
                    marker.points.append(p)
            # marker_array.markers.append(marker)
            # id = 0
            # for m in marker_array.markers:
            #     m.id = id
            #     id += 1
            self.mk_pub.publish(marker)
            # marker_array.markers = []
            
    def viz_car(self, labels, bbox3d, frame_id):
        bbox_car = []
        idx = 0
        for cls in labels:
            if cls == ClassName.car.value:
                bbox_car.append(bbox3d[idx])
            idx = idx + 1
        car_array = np.array(bbox_car)
        self.viz(car_array, frame_id)
    
    def viz_person(self, labels, bbox3d, frame_id):
        bbox_person = []
        idx = 0
        for cls in labels:
            if cls == ClassName.person.value:
                bbox_person.append(bbox3d[idx])
            idx = idx + 1
        person_array = np.array(bbox_person)
        self.viz(person_array, frame_id)
    
    def pub_trackbox(self, bbox3d, frame_id):
        
        if bbox3d.shape[0] == 0:
            print("no %s is detected"%category)
        else:

            # self.viz(bbox3d=bbox3d, frame_id=self.frame_id)
            corners3d = boxes_to_corners_3d(bbox3d)#(N,8,3)
            box_num = corners3d.shape[0]
            track_box = trackbox()

            xy = dict(x=0, y=4)
            for corners in corners3d:
                for k, v in xy.items():
                    xy_id = 1
                    for i in range(v, v+4):
                        attr = k + str(xy_id)
                        xy_id += 1
                        ebox = getattr(track_box, attr)
                        for j in range(3):
                            ebox.append(corners[i][j])

            track_box.header.frame_id = frame_id
            track_box.header.stamp = rospy.Time.now()
            track_box.box_num = box_num
            self.box_pub.publish(track_box)
            

            








